import os
import torch
import numpy as np
from torch.utils.data import Dataset,DataLoader
from hparams import hparams
import random

class Tacotron2_Dataset(Dataset):
    def __init__(self,para):
        self.file_scp = para.train_scp  # 训练集路径
        files = np.loadtxt(self.file_scp,dtype = 'str',delimiter = '|')
        # 编号|音素|音素对应的数字编号
        self.file_ids = files[:,0].tolist()
        self.index_phone = files[:,2].tolist()
        self.para = para
        
    # 读取音频特征
    def get_mel(self, file_id):
        file_fea = os.path.join(self.para.path_fea,file_id+'.npy')
        melspec = torch.from_numpy(np.load(file_fea))
        return melspec
    
    # 读取文本编码序列
    def get_text(self, str_phones):
        phone_ids = [int(id) for id in str_phones.split()]
        return torch.IntTensor(phone_ids)
    
    # 获取 文本/特征 对
    def get_mel_text_pair(self, file_id,str_phones_ids):
        text = self.get_text(str_phones_ids)
        mel = self.get_mel(file_id)
        return (text, mel)
        
    def __getitem__(self, index):
        return self.get_mel_text_pair(self.file_ids[index],self.index_phone[index])
    
    def __len__(self):
        return len(self.file_ids)

# 特征处理    
class TextMelCollate():
    """ 
        通过补0的方法使一个 batch 内 的  text(输入) 和  mel(目标) 一样长
        对 mel 进行补0的时候 要让最长的mel 是 frames per step 的整数倍
    """
    def __init__(self, n_frames_per_step):
        self.n_frames_per_step = n_frames_per_step
    def __call__(self, batch):
        # Right zero-pad all one-hot text sequences to max input length
        input_lengths, ids_sorted_decreasing = torch.sort(
            torch.LongTensor([len(x[0]) for x in batch]),# x[0] 文本
            dim=0, descending=True)# 降序
        max_input_len = input_lengths[0] # 最大长度

        text_padded = torch.LongTensor(len(batch), max_input_len)
        text_padded.zero_() # 确保所有未使用的位置都被填充为0
        for i in range(len(ids_sorted_decreasing)):
            text = batch[ids_sorted_decreasing[i]][0] # 文本特征
            # 将获取的文本特征text复制到text_padded张量的第i行，从第0列开始，直到text的长度
            text_padded[i, :text.size(0)] = text

        # 处理音频特征：Right zero-pad mel-spec
        num_mels = batch[0][1].size(0)
        max_target_len = max([x[1].size(1) for x in batch]) # 帧长
        if max_target_len % self.n_frames_per_step != 0:  #保证特征的帧长是n_frames_per_step的整数倍，每次处理的帧数
            max_target_len += self.n_frames_per_step - max_target_len % self.n_frames_per_step
            assert max_target_len % self.n_frames_per_step == 0

        # include mel padded and gate padded
        mel_padded = torch.FloatTensor(len(batch), num_mels, max_target_len)
        mel_padded.zero_()
        gate_padded = torch.FloatTensor(len(batch), max_target_len) # 构建gete 目标 判断生成什么时候结束
        gate_padded.zero_()
        output_lengths = torch.LongTensor(len(batch))
        for i in range(len(ids_sorted_decreasing)):# 顺序与文本特征对应
            mel = batch[ids_sorted_decreasing[i]][1]
            mel_padded[i, :, :mel.size(1)] = mel # 进行数据填充
            gate_padded[i, mel.size(1)-1:] = 1 # 最后一帧往后都是1
            output_lengths[i] = mel.size(1)
        # 文本特征、文本特征长度、音频特征、gate特征、音频特征长度
        return text_padded, input_lengths, mel_padded, gate_padded, \
            output_lengths